"""For basic init and call"""
import os

from langchain.agents import create_sql_agent
from langchain.agents.agent_toolkits import SQLDatabaseToolkit
from langchain.llms import QianfanLLMEndpoint
from langchain.sql_database import SQLDatabase

os.environ["QIANFAN_AK"] = "ZGgmCMM95MfabqYLG2swwRWM"
os.environ["QIANFAN_SK"] = "yDAIOpvEu7vm2phjYDT5sdIg5h4LBb4v"

llm = QianfanLLMEndpoint(model="ERNIE-Bot", streaming=True)

# 数据库
db_user = "root"
db_password = "root123"
db_host = "localhost:3306"
db_name = "group_insurance"
db = SQLDatabase.from_uri(f"mysql+pymysql://{db_user}:{db_password}@{db_host}/{db_name}")

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    verbose=True
)

agent_executor.run("列出订单数量最多的前 3 个客户名称")
